# @Time    : 2019/1/14 21:00
# @Author  : heyin
import pandas as pd
import numpy as np
import pandas_datareader.data as web


def boll_bands(data, ndays):
    """
    计算布林带
    :param data: 股票的df格式数据
    :param ndays: 计算使用的简单移动均线周期
    :return:
    """
    ma = pd.Series(data['Close'].rolling(ndays).mean().round(2))  # 计算nday均线
    # pandas.std() 默认是除以n-1 的，即是无偏的，如果想和numpy.std() 一样有偏，需要加上参数ddof=0
    # 此处添加ddof的原因是wind和yahoo的计算均采用的有偏值进行的计算
    std = pd.Series(data['Close'].rolling(ndays).std(ddof=0).round(2))  # 计算nday标准差，有偏
    b1 = ma + (2 * std)  # 此处的2就是Standard Deviations 标准偏差
    B1 = pd.Series(b1, name='UpperBollingerBand')  # 下边join的时候，这里的name变为列标签，没有name
    # data = data.join(ma)  # 上边不写name 这里报错
    data = data.join(B1)

    b2 = ma - (2 * std)
    B2 = pd.Series(b2, name='LowerBollingerBand')
    data = data.join(B2)
    return data


def sma(data, ndays):
    """
    简单移动平均
    :param data: 股票的df格式数据
    :param ndays: 计算使用的简单移动均线周期
    :return:
    """
    _sma = pd.Series(data['Close'].rolling(ndays).mean().round(2), name='SMA_%s' % ndays)
    data = data.join(_sma)
    return data


def ewma(data, ndays):
    """
    指数加权滑动平均
    :param data: 股票的df格式数据
    :param ndays: 计算使用的周期
    :return:
    """
    # pandas 指数加权滑动（ewm）, 指数加权滑动平均（ewma），在前者之上取均值
    _ema = pd.Series(data['Close'].ewm(span=ndays, min_periods=ndays - 1).mean().round(2), name='EWMA_%s' % ndays)
    data = data.join(_ema)
    return data


# Commodity Channel Index  商品通道指标，CCI指标又叫顺势指标
def cci(data, ndays):
    """
    Commodity Channel Index  商品通道指标，CCI指标又叫顺势指标，计算方法参考stockstats库实现
    :param data: 股票的df格式数据
    :param ndays: 计算使用的周期
    :return:
    """
    TP = (data['High'] + data['Low'] + data['Close']) / 3
    # np.fabs 绝对值，注意这里0.015后的参数，不是标准差，如此求得结果才与yahoo数据相同
    CCI = pd.Series(
        (TP - TP.rolling(ndays).mean()) / (
            0.015 * TP.rolling(ndays).apply(lambda x: np.fabs(x - x.mean()).mean(), raw=False)),
        name='CCI_%s' % ndays)

    data = data.join(CCI.round(2))
    return data


def macd(data, short=12, long=26, mid=9):
    """
    计算MACD
    :param data: df数据
    :param short: 短期
    :param long: 长期
    :param mid:
    :return:
    """
    # 计算短期和长期的ewma，使用pandas的ewm得到指数加权的方法
    short_ewma = pd.Series(data['Close']).ewm(span=short).mean()
    long_ewma = pd.Series(data['Close']).ewm(span=long).mean()
    # 计算macd的三项指标
    data['DIFF'] = (short_ewma - long_ewma).round(2)
    data['DEA'] = pd.Series(data['DIFF']).ewm(span=mid).mean().round(2)
    data['MACD'] = 2 * (data['DIFF'] - data['DEA']).round(2)
    return data


def kdj(data, n=9, m=3):
    """
    计算kdj
    :param data:
    :param n:
    :param m:
    :return:
    """
    low_list = data['Low'].rolling(n).min()
    high_list = data['High'].rolling(n).max()
    rsv = (data['Close'] - low_list) / (high_list - low_list) * 100
    data['KDJ_K'] = rsv.ewm(alpha=1 / m, adjust=False).mean()
    data['KDJ_D'] = data['KDJ_K'].ewm(alpha=1 / m, adjust=False).mean()
    data['KDJ_J'] = 3 * data['KDJ_K'] - 2 * data['KDJ_D']
    return data


if __name__ == '__main__':
    data = web.DataReader('^IXIC', data_source='yahoo', start='2/5/1971', end='1/14/2019')
    n = 20

    # 测试布林带
    # NIFTY_BBANDS = boll_bands(data, n)
    # print(NIFTY_BBANDS.loc[:, ['LowerBollingerBand', 'UpperBollingerBand']])
    # print(NIFTY_BBANDS.columns)

    # 测试sma
    # data = sma(data, n)
    # print(data.loc[:, ['SMA_%s' % n]])

    # 测试ewma
    # data = ewma(data, n)
    # print(data.loc[:, ["EWMA_%s" % n]])

    # 测试cci
    # data = cci(data, n)
    # print(data.loc[:, ["CCI_%s" % n]])

    # data = macd(data)
    # print(data.loc[:, ['diff', 'macd', 'dea']])

    data = kdj(data)
    print(data.loc[:, ['KDJ_K', 'KDJ_D', 'KDJ_J']])

    # TODO 需要知道ewm的参数是什么意思，如span，adjust，alpha